// Redis-based distributed storage backend
use codegraph_core::{
    Result, CodeGraphError, GraphStorage, Transaction, GraphQuery, QueryResult,
    CodeNode, CodeEdge, NodeId, EdgeId, EdgeType,
};
use redis::{Client, AsyncCommands, RedisResult};
use serde::{Serialize, Deserialize};
use std::sync::Arc;
use std::collections::HashMap;
use tokio::sync::RwLock;
use uuid::Uuid;
use async_trait::async_trait;

/// Helper function to create storage errors
fn storage_error(message: String) -> CodeGraphError {
    CodeGraphError::Storage {
        code: codegraph_core::ErrorCode::StorageError,
        message,
        context: None,
    }
}

/// Helper function to create configuration errors
fn config_error(message: String) -> CodeGraphError {
    CodeGraphError::Configuration {
        message,
        config_path: None,
    }
}

/// Redis storage configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RedisConfig {
    /// Redis connection URL
    pub url: String,
    /// Connection pool size
    pub pool_size: usize,
    /// Key prefix for namespacing
    pub key_prefix: String,
    /// TTL for cached data (in seconds)
    pub ttl: Option<u64>,
    /// Enable clustering
    pub cluster_mode: bool,
    /// Cluster nodes (if cluster_mode is true)
    pub cluster_nodes: Vec<String>,
    /// Connection timeout in milliseconds
    pub connection_timeout_ms: u64,
    /// Command timeout in milliseconds
    pub command_timeout_ms: u64,
}

impl Default for RedisConfig {
    fn default() -> Self {
        Self {
            url: "redis://127.0.0.1:6379".to_string(),
            pool_size: 10,
            key_prefix: "codegraph:".to_string(),
            ttl: None,
            cluster_mode: false,
            cluster_nodes: Vec::new(),
            connection_timeout_ms: 5000,
            command_timeout_ms: 10000,
        }
    }
}

/// Redis storage backend implementation
pub struct RedisStorage {
    client: Client,
    config: RedisConfig,
    connection_pool: Arc<RwLock<Vec<redis::aio::Connection>>>,
}

impl std::fmt::Debug for RedisStorage {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("RedisStorage")
            .field("config", &self.config)
            .field("connection_pool_size", &"<connection_pool>")
            .finish()
    }
}

/// Redis key types for different data structures
#[derive(Debug, Clone)]
enum RedisKey {
    Node(Uuid),
    Edge(Uuid),
    Graph(String),
    Index(String),
    Metadata(String),
    Lock(String),
}

impl RedisKey {
    fn to_string(&self, prefix: &str) -> String {
        match self {
            RedisKey::Node(id) => format!("{}node:{}", prefix, id),
            RedisKey::Edge(id) => format!("{}edge:{}", prefix, id),
            RedisKey::Graph(name) => format!("{}graph:{}", prefix, name),
            RedisKey::Index(name) => format!("{}index:{}", prefix, name),
            RedisKey::Metadata(key) => format!("{}meta:{}", prefix, key),
            RedisKey::Lock(resource) => format!("{}lock:{}", prefix, resource),
        }
    }
}

/// Distributed lock implementation
pub struct DistributedLock {
    key: String,
    value: String,
    ttl: u64,
    client: Client,
}

impl DistributedLock {
    pub async fn acquire(
        client: &Client,
        resource: &str,
        ttl_seconds: u64,
    ) -> Result<Option<Self>> {
        let mut conn = client.get_async_connection().await
            .map_err(|e| storage_error(e.to_string()))?;
        
        let lock_key = format!("lock:{}", resource);
        let lock_value = Uuid::new_v4().to_string();
        
        // Try to acquire lock with SET NX EX using Lua script
        let script = r#"
            if redis.call("SET", KEYS[1], ARGV[1], "NX", "EX", ARGV[2]) then
                return "OK"
            else
                return nil
            end
        "#;

        let result: RedisResult<Option<String>> = redis::Script::new(script)
            .key(&lock_key)
            .arg(&lock_value)
            .arg(ttl_seconds)
            .invoke_async(&mut conn)
            .await;
        
        match result {
            Ok(Some(_)) => Ok(Some(DistributedLock {
                key: lock_key,
                value: lock_value,
                ttl: ttl_seconds,
                client: client.clone(),
            })),
            Ok(None) => Ok(None), // Lock already held
            Err(_) => Ok(None), // Error acquiring lock
        }
    }
    
    pub async fn release(&self) -> Result<bool> {
        let mut conn = self.client.get_async_connection().await
            .map_err(|e| storage_error(e.to_string()))?;
        
        // Use Lua script to ensure atomic release
        let script = r#"
            if redis.call("GET", KEYS[1]) == ARGV[1] then
                return redis.call("DEL", KEYS[1])
            else
                return 0
            end
        "#;
        
        let result: RedisResult<i32> = redis::Script::new(script)
            .key(&self.key)
            .arg(&self.value)
            .invoke_async(&mut conn)
            .await;
        
        match result {
            Ok(1) => Ok(true),
            Ok(0) => Ok(false),
            Ok(_) => Ok(false), // Unexpected return value
            Err(e) => Err(storage_error(e.to_string())),
        }
    }
    
    pub async fn extend(&mut self, additional_ttl: u64) -> Result<bool> {
        let mut conn = self.client.get_async_connection().await
            .map_err(|e| storage_error(e.to_string()))?;
        
        // Use Lua script to extend TTL atomically
        let script = r#"
            if redis.call("GET", KEYS[1]) == ARGV[1] then
                return redis.call("EXPIRE", KEYS[1], ARGV[2])
            else
                return 0
            end
        "#;
        
        let result: RedisResult<i32> = redis::Script::new(script)
            .key(&self.key)
            .arg(&self.value)
            .arg(self.ttl + additional_ttl)
            .invoke_async(&mut conn)
            .await;
        
        match result {
            Ok(1) => {
                self.ttl += additional_ttl;
                Ok(true)
            },
            Ok(0) => Ok(false),
            Ok(_) => Ok(false), // Unexpected return value
            Err(e) => Err(storage_error(e.to_string())),
        }
    }
}

impl RedisStorage {
    /// Create a new Redis storage backend
    pub async fn new(config: RedisConfig) -> Result<Self> {
        let client = if config.cluster_mode {
            // TODO: Implement cluster client
            return Err(config_error(
                "Redis cluster mode not yet implemented".to_string()
            ));
        } else {
            Client::open(config.url.clone())
                .map_err(|e| storage_error(e.to_string()))?
        };
        
        // Test connection
        let mut conn = client.get_async_connection().await
            .map_err(|e| storage_error(e.to_string()))?;
        
        // Ping to verify connection
        let _: String = redis::cmd("PING").query_async(&mut conn).await
            .map_err(|e| storage_error(e.to_string()))?;
        
        Ok(Self {
            client,
            config,
            connection_pool: Arc::new(RwLock::new(Vec::new())),
        })
    }
    
    /// Get a connection from the pool or create a new one
    async fn get_connection(&self) -> Result<redis::aio::Connection> {
        // Try to get from pool first
        {
            let mut pool = self.connection_pool.write().await;
            if let Some(conn) = pool.pop() {
                return Ok(conn);
            }
        }
        
        // Create new connection
        self.client.get_async_connection().await
            .map_err(|e| storage_error(e.to_string()))
    }
    
    /// Return a connection to the pool
    async fn return_connection(&self, conn: redis::aio::Connection) {
        let mut pool = self.connection_pool.write().await;
        if pool.len() < self.config.pool_size {
            pool.push(conn);
        }
        // If pool is full, connection will be dropped
    }
    
    /// Execute a Redis command with automatic retry
    async fn execute_with_retry<T, F, Fut>(&self, operation: F) -> Result<T>
    where
        F: Fn(redis::aio::Connection) -> Fut,
        Fut: std::future::Future<Output = RedisResult<(T, redis::aio::Connection)>>,
    {
        const MAX_RETRIES: usize = 3;
        let mut last_error = None;
        
        for attempt in 0..MAX_RETRIES {
            let conn = self.get_connection().await?;
            
            match operation(conn).await {
                Ok((result, conn)) => {
                    self.return_connection(conn).await;
                    return Ok(result);
                }
                Err(e) => {
                    last_error = Some(e);
                    if attempt < MAX_RETRIES - 1 {
                        // Wait before retry
                        tokio::time::sleep(tokio::time::Duration::from_millis(100 * (attempt + 1) as u64)).await;
                    }
                }
            }
        }
        
        Err(storage_error(
            last_error.unwrap().to_string()
        ))
    }
    
    /// Acquire a distributed lock
    pub async fn acquire_lock(&self, resource: &str, ttl_seconds: u64) -> Result<Option<DistributedLock>> {
        DistributedLock::acquire(&self.client, resource, ttl_seconds).await
    }
    
    /// Store data with optional TTL
    async fn store_with_ttl(&self, key: &str, value: &[u8], ttl: Option<u64>) -> Result<()> {
        self.execute_with_retry(|mut conn| async move {
            let result = if let Some(ttl_seconds) = ttl {
                conn.set_ex(key, value, ttl_seconds).await
            } else {
                conn.set(key, value).await
            };
            
            result.map(|r| (r, conn))
        }).await
    }
    
    /// Get data by key
    async fn get_data(&self, key: &str) -> Result<Option<Vec<u8>>> {
        self.execute_with_retry(|mut conn| async move {
            let result: RedisResult<Option<Vec<u8>>> = conn.get(key).await;
            result.map(|r| (r, conn))
        }).await
    }
    
    /// Delete data by key
    async fn delete_data(&self, key: &str) -> Result<bool> {
        self.execute_with_retry(|mut conn| async move {
            let result: RedisResult<i32> = conn.del(key).await;
            result.map(|r| (r > 0, conn))
        }).await
    }
    
    /// Check if key exists
    pub async fn exists(&self, key: &str) -> Result<bool> {
        self.execute_with_retry(|mut conn| async move {
            let result: RedisResult<bool> = conn.exists(key).await;
            result.map(|r| (r, conn))
        }).await
    }
    
    /// Get multiple keys at once
    async fn get_multiple(&self, keys: &[String]) -> Result<Vec<Option<Vec<u8>>>> {
        if keys.is_empty() {
            return Ok(Vec::new());
        }
        
        self.execute_with_retry(|mut conn| async move {
            let result: RedisResult<Vec<Option<Vec<u8>>>> = conn.get(keys).await;
            result.map(|r| (r, conn))
        }).await
    }
    
    /// Set multiple key-value pairs atomically
    async fn set_multiple(&self, pairs: &[(String, Vec<u8>)]) -> Result<()> {
        if pairs.is_empty() {
            return Ok(());
        }
        
        self.execute_with_retry(|mut conn| async move {
            let mut pipe = redis::pipe();
            for (key, value) in pairs {
                pipe.set(key, value);
            }
            
            let result: RedisResult<()> = pipe.query_async(&mut conn).await;
            result.map(|r| (r, conn))
        }).await
    }
    
    /// Add to a set
    async fn add_to_set(&self, set_key: &str, member: &str) -> Result<bool> {
        self.execute_with_retry(|mut conn| async move {
            let result: RedisResult<i32> = conn.sadd(set_key, member).await;
            result.map(|r| (r > 0, conn))
        }).await
    }
    
    /// Get all members of a set
    async fn get_set_members(&self, set_key: &str) -> Result<Vec<String>> {
        self.execute_with_retry(|mut conn| async move {
            let result: RedisResult<Vec<String>> = conn.smembers(set_key).await;
            result.map(|r| (r, conn))
        }).await
    }
    
    /// Remove from a set
    async fn remove_from_set(&self, set_key: &str, member: &str) -> Result<bool> {
        self.execute_with_retry(|mut conn| async move {
            let result: RedisResult<i32> = conn.srem(set_key, member).await;
            result.map(|r| (r > 0, conn))
        }).await
    }
}

impl RedisStorage {
    /// Store data with key
    pub async fn store(&mut self, key: &str, data: &[u8]) -> Result<()> {
        let redis_key = format!("{}{}", self.config.key_prefix, key);
        self.store_with_ttl(&redis_key, data, self.config.ttl).await
    }

    /// Retrieve data by key
    pub async fn retrieve(&self, key: &str) -> Result<Option<Vec<u8>>> {
        let redis_key = format!("{}{}", self.config.key_prefix, key);
        self.get_data(&redis_key).await
    }

    /// Delete data by key
    pub async fn delete(&mut self, key: &str) -> Result<bool> {
        let redis_key = format!("{}{}", self.config.key_prefix, key);
        self.delete_data(&redis_key).await
    }

    /// Check if key exists
    pub async fn key_exists(&self, key: &str) -> Result<bool> {
        let redis_key = format!("{}{}", self.config.key_prefix, key);
        self.exists(&redis_key).await
    }

    /// List all keys with optional prefix
    pub async fn list_keys(&self, prefix: Option<&str>) -> Result<Vec<String>> {
        let pattern = match prefix {
            Some(p) => format!("{}{}*", self.config.key_prefix, p),
            None => format!("{}*", self.config.key_prefix),
        };

        // Execute keys command directly without retry wrapper for now
        let mut conn = self.get_connection().await?;
        let keys: Vec<String> = conn.keys(&pattern).await
            .map_err(|e| storage_error(e.to_string()))?;
        self.return_connection(conn).await;

        // Remove prefix from keys
        let prefix_len = self.config.key_prefix.len();
        let result = keys.into_iter()
            .map(|k| k[prefix_len..].to_string())
            .collect();

        Ok(result)
    }

    /// Clear all data
    pub async fn clear(&mut self) -> Result<()> {
        // Get all keys with our prefix
        let keys = self.list_keys(None).await?;

        if !keys.is_empty() {
            let redis_keys: Vec<String> = keys.into_iter()
                .map(|k| format!("{}{}", self.config.key_prefix, k))
                .collect();

            // Execute delete command directly
            let mut conn = self.get_connection().await?;
            let _: i32 = conn.del(&redis_keys).await
                .map_err(|e| storage_error(e.to_string()))?;
            self.return_connection(conn).await;
        }

        Ok(())
    }
}

// Implement GraphStorage trait for RedisStorage
#[async_trait]
impl GraphStorage for RedisStorage {
    async fn save_node(&self, node: &dyn CodeNode) -> Result<()> {
        let key = RedisKey::Node(node.id()).to_string(&self.config.key_prefix);

        // For now, use a simple binary format
        // In a real implementation, we'd use a proper serialization format
        let mut data = Vec::new();

        // Serialize ID
        let id_bytes = node.id().to_string().into_bytes();
        data.extend_from_slice(&(id_bytes.len() as u32).to_le_bytes());
        data.extend_from_slice(&id_bytes);

        // Serialize type
        let type_bytes = node.node_type().to_string().into_bytes();
        data.extend_from_slice(&(type_bytes.len() as u32).to_le_bytes());
        data.extend_from_slice(&type_bytes);

        // Serialize name
        let name_bytes = node.name().as_bytes();
        data.extend_from_slice(&(name_bytes.len() as u32).to_le_bytes());
        data.extend_from_slice(name_bytes);

        // Store the serialized data
        self.store_with_ttl(&key, &data, self.config.ttl).await
    }

    async fn get_node(&self, id: NodeId) -> Result<Option<Box<dyn CodeNode>>> {
        let key = RedisKey::Node(id).to_string(&self.config.key_prefix);
        if let Some(_data) = self.get_data(&key).await? {
            // TODO: Deserialize to concrete node type
            // For now, return None as we need proper deserialization
            Ok(None)
        } else {
            Ok(None)
        }
    }

    async fn delete_node(&self, id: NodeId) -> Result<bool> {
        let key = RedisKey::Node(id).to_string(&self.config.key_prefix);
        self.delete_data(&key).await
    }

    async fn save_edge(&self, edge: &dyn CodeEdge) -> Result<()> {
        let key = RedisKey::Edge(edge.id()).to_string(&self.config.key_prefix);

        // Simple binary serialization for edges
        let mut data = Vec::new();

        // Serialize ID
        let id_bytes = edge.id().to_string().into_bytes();
        data.extend_from_slice(&(id_bytes.len() as u32).to_le_bytes());
        data.extend_from_slice(&id_bytes);

        // Serialize source
        let source_bytes = edge.source().to_string().into_bytes();
        data.extend_from_slice(&(source_bytes.len() as u32).to_le_bytes());
        data.extend_from_slice(&source_bytes);

        // Serialize target
        let target_bytes = edge.target().to_string().into_bytes();
        data.extend_from_slice(&(target_bytes.len() as u32).to_le_bytes());
        data.extend_from_slice(&target_bytes);

        // Serialize edge type
        let type_bytes = format!("{:?}", edge.edge_type()).into_bytes();
        data.extend_from_slice(&(type_bytes.len() as u32).to_le_bytes());
        data.extend_from_slice(&type_bytes);

        self.store_with_ttl(&key, &data, self.config.ttl).await
    }

    async fn get_edges(&self, _node_id: NodeId, _edge_type: Option<&EdgeType>) -> Result<Vec<Box<dyn CodeEdge>>> {
        // TODO: Implement edge retrieval with proper indexing
        // For now, return empty vector
        Ok(Vec::new())
    }

    async fn delete_edge(&self, id: EdgeId) -> Result<bool> {
        let key = RedisKey::Edge(id).to_string(&self.config.key_prefix);
        self.delete_data(&key).await
    }

    async fn save_batch(
        &self,
        nodes: &[&dyn CodeNode],
        edges: &[&dyn CodeEdge],
    ) -> Result<()> {
        // Use Redis pipeline for batch operations
        let mut pipeline = redis::pipe();

        // Add nodes to pipeline
        for node in nodes {
            let key = RedisKey::Node(node.id()).to_string(&self.config.key_prefix);

            // Use the same serialization as save_node
            let mut data = Vec::new();
            let id_bytes = node.id().to_string().into_bytes();
            data.extend_from_slice(&(id_bytes.len() as u32).to_le_bytes());
            data.extend_from_slice(&id_bytes);

            let type_bytes = node.node_type().to_string().into_bytes();
            data.extend_from_slice(&(type_bytes.len() as u32).to_le_bytes());
            data.extend_from_slice(&type_bytes);

            let name_bytes = node.name().as_bytes();
            data.extend_from_slice(&(name_bytes.len() as u32).to_le_bytes());
            data.extend_from_slice(name_bytes);

            if let Some(ttl) = self.config.ttl {
                pipeline.set_ex(&key, &data, ttl);
            } else {
                pipeline.set(&key, &data);
            }
        }

        // Add edges to pipeline
        for edge in edges {
            let key = RedisKey::Edge(edge.id()).to_string(&self.config.key_prefix);

            // Use the same serialization as save_edge
            let mut data = Vec::new();
            let id_bytes = edge.id().to_string().into_bytes();
            data.extend_from_slice(&(id_bytes.len() as u32).to_le_bytes());
            data.extend_from_slice(&id_bytes);

            let source_bytes = edge.source().to_string().into_bytes();
            data.extend_from_slice(&(source_bytes.len() as u32).to_le_bytes());
            data.extend_from_slice(&source_bytes);

            let target_bytes = edge.target().to_string().into_bytes();
            data.extend_from_slice(&(target_bytes.len() as u32).to_le_bytes());
            data.extend_from_slice(&target_bytes);

            let type_bytes = format!("{:?}", edge.edge_type()).into_bytes();
            data.extend_from_slice(&(type_bytes.len() as u32).to_le_bytes());
            data.extend_from_slice(&type_bytes);

            if let Some(ttl) = self.config.ttl {
                pipeline.set_ex(&key, &data, ttl);
            } else {
                pipeline.set(&key, &data);
            }
        }

        // Execute pipeline directly
        let mut conn = self.get_connection().await?;
        let _: () = pipeline.query_async(&mut conn).await
            .map_err(|e| CodeGraphError::storage_error(e.to_string()))?;
        self.return_connection(conn).await;

        Ok(())
    }

    async fn query(&self, _query: &GraphQuery) -> Result<QueryResult> {
        // TODO: Implement distributed query execution
        // For now, return empty result
        Ok(QueryResult {
            nodes: Vec::new(),
            edges: Vec::new(),
            metadata: HashMap::new(),
            execution_time_ms: 0,
        })
    }

    async fn begin_transaction(&self) -> Result<Box<dyn Transaction>> {
        Ok(Box::new(RedisTransaction::new(self.client.clone(), self.config.clone())))
    }
}

/// Redis transaction implementation
pub struct RedisTransaction {
    client: Client,
    config: RedisConfig,
    operations: Vec<TransactionOperation>,
}

#[derive(Debug)]
enum TransactionOperation {
    SaveNode { key: String, data: Vec<u8> },
    SaveEdge { key: String, data: Vec<u8> },
    DeleteNode { key: String },
    DeleteEdge { key: String },
}

impl RedisTransaction {
    pub fn new(client: Client, config: RedisConfig) -> Self {
        Self {
            client,
            config,
            operations: Vec::new(),
        }
    }
}

#[async_trait]
impl Transaction for RedisTransaction {
    async fn commit(self: Box<Self>) -> Result<()> {
        if self.operations.is_empty() {
            return Ok(());
        }

        // Execute all operations in a Redis transaction
        let mut conn = self.client.get_async_connection().await
            .map_err(|e| CodeGraphError::storage_error(e.to_string()))?;

        // Start Redis transaction
        let mut pipe = redis::pipe();
        pipe.atomic();

        // Add all operations to the pipeline
        for op in &self.operations {
            match op {
                TransactionOperation::SaveNode { key, data } |
                TransactionOperation::SaveEdge { key, data } => {
                    if let Some(ttl) = self.config.ttl {
                        pipe.set_ex(key, data, ttl);
                    } else {
                        pipe.set(key, data);
                    }
                }
                TransactionOperation::DeleteNode { key } |
                TransactionOperation::DeleteEdge { key } => {
                    pipe.del(key);
                }
            }
        }

        // Execute the transaction
        let _: () = pipe.query_async(&mut conn).await
            .map_err(|e| CodeGraphError::storage_error(e.to_string()))?;

        Ok(())
    }

    async fn rollback(self: Box<Self>) -> Result<()> {
        // Redis transactions are atomic, so rollback is automatic on failure
        // Just clear the operations
        Ok(())
    }

    async fn save_node(&self, _node: &dyn CodeNode) -> Result<()> {
        // For now, just return Ok since we can't modify self in this trait method
        // In a real implementation, we'd need to use interior mutability
        Ok(())
    }

    async fn save_edge(&self, _edge: &dyn CodeEdge) -> Result<()> {
        // For now, just return Ok since we can't modify self in this trait method
        // In a real implementation, we'd need to use interior mutability
        Ok(())
    }

    async fn delete_node(&self, _id: NodeId) -> Result<bool> {
        // For now, just return Ok since we can't modify self in this trait method
        // In a real implementation, we'd need to use interior mutability
        Ok(false)
    }

    async fn delete_edge(&self, _id: EdgeId) -> Result<bool> {
        // For now, just return Ok since we can't modify self in this trait method
        // In a real implementation, we'd need to use interior mutability
        Ok(false)
    }
}
